#!/usr/bin/python
"""Alignment is used for representing a bunch of binding sites.

Alignment is supposed to use as an instance which is initialised by a file having a bunch of binding
sites which bind to particular protein. These all binding sites must be of same length.
"""

import sys

class EmptyAlignmentFileException(Exception):
    "This will be raised when alignment contains no sequences."
    def __init__(self):
        Exception.__init__(self)
        print "Alignment file is empty and contains no sequences."

class Alignment:
    
    sequences = []
    __b = -1
    __w = -1
    # temp_sequences is used for taking care of blank lines and not used anywhere else
    temp_sequences = []

    def __init__(self, filename):
        try:
            file = open(filename, 'r')
        except:
            print 'oops..some error. check filename'
            sys.exit(1)

        self.temp_sequences = file.readlines()
        
        for i in xrange(len(self.temp_sequences)):
            self.temp_sequences[i] = (self.temp_sequences[i]).strip()

            if len(self.temp_sequences[i]) != 0:
                (self.sequences).append(self.temp_sequences[i])
                self.sequences[-1] = (self.sequences[-1]).lower()

        del self.temp_sequences
        file.close() 
    
    # Property which returns total number of sequences in given Alignment file       
    def __ReturnTotal(self):
        return len(self.sequences)
    
    # Property which returns length of each sequence in given alignment file
    def __ReturnLength(self):
        if self.Total == 0:
            raise EmptyAlignmentFileException
        else:
            return len(self.sequences[0])
    
    # Property which returns index of sequence at maximum distance from all other ones
    def __ReturnWorst(self):
        """Get the index of least similar sequence in given alignment."""
        if self.__w != -1:
            return self.__w
        # index of sequence at maximum distance from all sequences in given alignment file
        idx = 0 
        total = self.Total
        length = self.Length
        freq_matrix = self.freq_matrix()
        
        if total == 0:
            # Alignment file contains no sequence
            raise EmptyAlignmentFileException
        
        min_prob = 1
        
        for seq_idx in xrange(total):
            seq_prob = 1
            seq = self.sequences[seq_idx]
            
            for base_idx in xrange(len(seq)):
                if seq[base_idx] == 'a':
                    seq_prob *= freq_matrix[base_idx]
                elif seq[base_idx] == 'c':
                    seq_prob *= freq_matrix[length + base_idx]
                elif seq[base_idx] == 'g':
                    seq_prob *= freq_matrix[2*length + base_idx]
                elif seq[base_idx] == 't':
                    seq_prob *= freq_matrix[3*length + base_idx]
                else:
                    print "some error in sequence data"
                    sys.exit(1)
                
            if seq_prob < min_prob:
                min_prob = seq_prob
                idx =  seq_idx
                
        self.__w = idx
        
        return idx
                
    # Property which returns index of sequence at maximum distance from all other ones
    def __ReturnBest(self):
        """Get the index of most similar sequence in given alignment."""
        if self.__b != -1:
            return self.__b
        # index of sequence at maximum distance from all sequences in given alignment file
        idx = 0
        total = self.Total
        length = self.Length
        freq_matrix = self.freq_matrix()

        if total == 0:
            # Alignment file contains no sequence
            raise EmptyAlignmentFileException

        max_prob = 0

        for seq_idx in xrange(total):
            seq_prob = 1
            seq = self.sequences[seq_idx]

            for base_idx in xrange(len(seq)):
                if seq[base_idx] == 'a':
                    seq_prob *= freq_matrix[base_idx]
                elif seq[base_idx] == 'c':
                    seq_prob *= freq_matrix[length + base_idx]
                elif seq[base_idx] == 'g':
                    seq_prob *= freq_matrix[2*length + base_idx]
                elif seq[base_idx] == 't':
                    seq_prob *= freq_matrix[3*length + base_idx]
                else:
                    print "some error in sequence data"
                    sys.exit(1)

            if seq_prob > max_prob:
                max_prob = seq_prob
                idx =  seq_idx

        self.__b = idx

        return idx
    
    def __ReturnConsensus(self):
        con = ''
        freq_matrix = self.freq_matrix()
        
        for i in xrange(self.Length):
            cur_freq = freq_matrix[i]
            base = 'a'
            if freq_matrix[i+self.Length] > cur_freq:
                base = 'c'
                cur_freq = freq_matrix[i+self.Length]
            if freq_matrix[i+2*(self.Length)] > cur_freq:
                base = 'g'
                cur_freq = freq_matrix[i+2*(self.Length)]
            if freq_matrix[i+3*(self.Length)] > cur_freq:
                base = 't'
                cur_freq = freq_matrix[i+3*(self.Length)]
                
            con += base
            
        return con

    # Total number of sequences in given Alignment
    Total = property(fget=__ReturnTotal, doc="Total number of sequences in given Alignment file.")
    
    # Index of sequence most similar to frequency matrix out of all given sequences.
    Best = property(fget=__ReturnBest, doc="Index of sequence at minimum distance from \
    all the other sequences in Alignment.")
    
    # Index of sequence at maximum distance from all the other sequences in given alignment file.
    Worst = property(fget=__ReturnWorst, doc="Index of sequence at maximum distance from \
    all the other sequences in Alignment.")
    
    # This property holds length of each sequence in given alignment file.
    Length = property(fget=__ReturnLength, doc="Length of each sequence in given Alignment file.")

    # This property holds consensus for given alignment based on base frequency at each position
    Consensus = property(fget=__ReturnConsensus, doc="Length of each sequence in given Alignment file.")
    
    def freq_matrix(self):
        "This returns the frequency matrix for given alignment file."
        length = self.Length
        matrix = [0 for tmp in xrange(length * 4)]
        
        for s in self.sequences:
            for i in xrange(len(s)):
                if s[i] == 'a':
                    matrix[i] += 1
                elif s[i] == 'c':
                    matrix[length+i] += 1
                elif s[i] == 'g':
                    matrix[length*2+i] += 1
                elif s[i] == 't':
                    matrix[length*3+i] += 1
                else:
                    print "sequence contains non primitive dna base. exiting."
                    sys.exit(1)
                    
        for i in xrange(len(matrix)):
            matrix[i] /= (float)(self.Total)
                    
        return matrix
    
    def print_freq_matrix(self):
        """Prints the frequency matrix for given alignment."""
        m = self.freq_matrix()

        for i in xrange(len(m)):
            if not i % self.Length:
                print
            print "%.3f" %(m[i]),
            
    def reset(self):
        """Reset best and worst sequence index.
        
        So that next time on their request they will be regenerated. This is
        important as things will change on addition and removal of sequences
        from alignment.
        """
        self.__b = -1
        self.__w = -1
        
    def add(self, seq):
        """Add given sequence to the given alignment."""
        if self.Total == 0:
            (self.sequences).append(seq)
            self.reset()
            return
        if len(seq) != self.Length:
            raise Exception, 'Attempt to add sequence of different length to alignment.'
        (self.sequences).append(seq)
        self.reset()
        
    def remove(self, index):
        """Remove sequence having index, index from the given alignment."""
        if index+1 > self.Total:
            raise Exception, 'Attempt to remove non-existing sequence from alignment.'
        del self.sequences[index]
        print "ok"
        self.reset()
        
    def __repr__(self):
        "Print representation for Alignment class."
        seqs = '\n'
        for seq in self.sequences:
            seqs += seq + '\n'
        return seqs
             
if __name__ == '__main__':
    a = Alignment(sys.argv[1])
    print a
